#include "rpc_stub.h"
#include "rpc_logger.h"
#include "object_pool.h"
#include "rpc_root.h"
#include <chrono>

namespace rpc {

Stub::~Stub() {
}

void StubCall::setData(TransportPtr transport, int size) {
    data_.reset(new char[size]);
    dataSize_ = size;
    transport->recv(data_.get(), size);
}

const char* StubCall::getData(int& size) {
    size = dataSize_;
    return data_.get();
}

MethodID StubCall::getMethodID() {
    return methodID_;
}

CallID StubCall::getCallID() {
    return callID_;
}

TransportPtr StubCall::getTransport() {
    return transport_;
}

ServiceID StubCall::getServiceID() {
    return serviceID_;
}

bool StubCall::isFinished() {
    return finished_;
}

void StubCall::setTransport(TransportPtr transport) {
    transport_ = transport;
}

void StubCall::setTimerHandle(TimerHandle timerHandle) {
    timerHandle_ = timerHandle;
}

TimerHandle StubCall::getTimerHandle() {
    return timerHandle_;
}

void StubCall::setCoroID(coroutine::CoroID coroID) {
    coroID_ = coroID;
}

coroutine::CoroID StubCall::getCoroID() {
    return coroID_;
}

void StubCall::setCallID(CallID callID) {
    callID_ = callID;
}

void StubCall::setMethodID(MethodID methodID) {
    methodID_ = methodID;
}

void StubCall::setServiceID(ServiceID serviceID) {
    serviceID_ = serviceID;
}

void StubCall::setStubCallID(StubCallID stubCallID) {
    stubCallID_ = stubCallID;
}

StubCallID StubCall::getStubCallID() {
    return stubCallID_;
}

void Stub::callCoro(StubCallPtr stubCall) {
    StubCallManagerRef.add(stubCall);
    auto stubCallID = stubCall->getStubCallID();
    auto coroID = coro_spawn([&](void* userData) {
        auto stubCallID = static_cast<StubCallID>(reinterpret_cast<std::uint64_t>(userData));
        auto stubCall = StubCallManagerRef.get(stubCallID);
        if (!stubCall) {
            return;
        }
        auto stub = StubManagerRef.getServiceByID(stubCall->getServiceID());
        if (!stub) {
            StubCallManagerRef.destroy(stubCallID);
            return;
        }
        if (stubCall->getTransport()->isClose()) {
            stub->decref();
            StubCallManagerRef.destroy(stubCallID);
            return;
        }
        try {
            stub->call(stubCall);
        } catch (std::exception& e) {
            Logbegin << "Call "
                << "["
                << stub->getServiceName()
                << ":"
                << stub->getMethodName(stubCall->getMethodID())
                << "] exception thrown:" << e.what()
                << Logend;
        }
        StubCallManagerRef.destroy(stubCallID);
        stub->decref();
    }, reinterpret_cast<void*>(stubCallID));
    if (coroID == coroutine::INVALID_CORO_ID) {
        Logbegin << "Call "
            << "["
            << getServiceName()
            << ":"
            << getMethodName(stubCall->getMethodID())
            << "] failed, start coroutine failed"
            << Logend;
        StubCallManagerRef.destroy(stubCallID);
        decref();
        return;
    }
    stubCall->setCoroID(coroID);
    stubCall->startExecTimer();
}

void Stub::callNormal(StubCallPtr stubCall) {
    StubCallManagerRef.add(stubCall);
    auto stubCallID = stubCall->getStubCallID();
    try {
        call(stubCall);
    } catch (std::exception& e) {
        Logbegin << "Call "
            << "["
            << getServiceName()
            << ":"
            << getMethodName(stubCall->getMethodID())
            << "] exception thrown:" << e.what()
            << Logend;
        StubCallManagerRef.destroy(stubCallID);
        decref();
        return;
    }
    if (stubCall->isFinished()) {
        StubCallManagerRef.destroy(stubCallID);
        decref();
        return;
    } else {
        stubCall->startExecTimer();
    }
}

void Stub::update(std::time_t ms) {
    tick(ms);
    if (factory_->isDynamic()) {
        if ((status_ == StubStatus::UPDATING) && zeroref()) {
            auto* dynamicFactory = dynamic_cast<StubFactoryDynamic*>(factory_);
            // Upgrade failure
            if (!reset(dynamicFactory->getEntryFunc(), &RpcRef)) {
                // Already rollback to old one
                status_ = StubStatus::RESOLVED;
            }
            // Avoid some branch donot change status to RESOLVED
            if (status_ != StubStatus::RESOLVED) {
                setStatus(StubStatus::RESOLVED);
                return;
            }
        }
    }
    if (callQueue_.empty() || !zeroref()) {
        return;
    }
    auto callPtr = callQueue_.back();
    callQueue_.pop_back();
    callStub(callPtr);
}

void Stub::callStub(TransportPtr transport, int bytes, const RpcCallRequestHeader& header) {
    auto callPtr = make_shared_ptr<StubCall>();
    callPtr->setTransport(transport);
    callPtr->setMethodID(header.methodID);
    callPtr->setCallID(header.callID);
    callPtr->setServiceID(serviceID_);
    callPtr->setData(transport, bytes);
    callStub(callPtr);
}

void Stub::callStub(StubCallPtr callPtr) {
    if (type_ != StubType::REENTRANT) {
        if (!zeroref()) {
            callQueue_.push_front(callPtr);
            return;
        }
    }
    incref();
    if (Singleton<Rpc>::instance()->isCoroMode()) {
        callCoro(callPtr);
    } else {
        callNormal(callPtr);
    }
}

void Stub::decref() {
    reference_-= 1;
}

void Stub::incref() {
    reference_ += 1;
}

bool Stub::zeroref() {
    return (reference_ == 0);
}

void Stub::setType(StubType type) {
    type_ = type;
}

StubType Stub::getType() {
    return type_;
}

ServiceID Stub::getServiceID() {
    return serviceID_;
}

void Stub::setServiceID(ServiceID serviceID) {
    serviceID_ = serviceID;
}

void Stub::setEntryFunc(BundleEntryFunc entryFunc) {
    entryFunc_ = entryFunc;
}

BundleEntryFunc Stub::getEntryFunc() {
    return entryFunc_;
}

void Stub::setStatus(StubStatus status) {
    status_ = status;
}

StubStatus Stub::getStatus() {
    return status_;
}

void Stub::setFactory(StubFactory* factory) {
    factory_ = factory;
}

StubFactory* Stub::getFactory() {
    return factory_;
}

void Stub::setContext(kratos::service::ServiceContext* ctx) {
    ctx_ = ctx;
}

kratos::service::ServiceContext* Stub::getContext() {
    return ctx_;
}

StubCall::~StubCall() {
    clean();
}

void StubCall::finish() {
    if (finished_) {
        return;
    }
    finished_ = true;
    StubCallManagerRef.destroy(stubCallID_);
}

bool StubCall::startExecTimer() {
    auto handle = StubManagerRef.getTimer().addTimerOnce(
    [&](std::uint64_t callID) {
        auto stubCall = StubCallManagerRef.get(callID);
        if (!stubCall) {
            return;
        }
        auto stub = StubManagerRef.getServiceByID(stubCall->getServiceID());
        if (stub) {
            Logbegin << "Service execution timeout "
                << "["
                << stub->getServiceName()
                << ":"
                << stub->getMethodName(stubCall->getMethodID())
                << "]"
                << Logend;
            stub->decref();
        }
        StubCallManagerRef.destroy(stubCall->getStubCallID());
    }, RPC_TIMEOUT, stubCallID_);
    timerHandle_ = handle;
    return (handle != nullptr);
}

void StubCall::stopExecTimer() {
    if (timerHandle_) {
        StubManagerRef.getTimer().cancelTimer(timerHandle_);
    }
    timerHandle_ = nullptr;
}

void StubCall::notFound() {
    RpcRef.notFound(transport_, callID_);
}

void StubCall::clean() {
    stopExecTimer();
    if (coroID_ != coroutine::INVALID_CORO_ID) {
        coro_close(coroID_);
        coroID_ = coroutine::INVALID_CORO_ID;
    }
    methodID_ = INVALID_METHOD_ID;
}

void buildHeader(RpcRetHeader& rheader, rpc::StubCallPtr stubCall, rpc::ServiceID serviceID, rpc::ErrorID errorID,  std::uint32_t bodySize) {
    rheader.header.length = sizeof(RpcRetHeader) + bodySize;
    rheader.header.type = static_cast<rpc::MsgTypeID>(rpc::RpcMsgType::RPC_RETURN);
    rheader.retHeader.serviceID = serviceID;
    rheader.retHeader.callID = stubCall->getCallID();
    rheader.retHeader.errorID = errorID;
    rheader.hton();
}

StubFactoryDynamic::~StubFactoryDynamic() {
    for (auto& ci : classList_) {
        unload(ci);
    }
}

StubManager::~StubManager() {
    for (auto& it : serviceMap_) {
        it.second->onBeforeDestory(&RpcRef);
    }
    serviceMap_.clear();
}

StubPtr StubManager::addService(StubPtr stub) {
    if (!stub->onAfterFork(&RpcRef)) {
        return nullptr;
    }
    stub->setServiceID(serviceID_++);
    typeMap_[stub->getUUID()].services.push_back(stub);
    serviceMap_.emplace(stub->getServiceID(), stub);
    return stub;
}

void StubManager::setStatus(ServiceUUID uuid, StubStatus status) {
    for (auto& it : serviceMap_) {
        if (it.second->getUUID() == uuid) {
            it.second->setStatus(status);
        }
    }
}

Timer& StubManager::getTimer() {
    return timer_;
}

StubPtr StubManager::getServiceByUUID(ServiceUUID uuid) {
    auto it = typeMap_.find(uuid);
    if (it == typeMap_.end()) {
        return nullptr;
    }
    if (it->second.services.empty()) {
        return nullptr;
    }
    if (it->second.index >= it->second.services.size()) {
        it->second.index = 0;
    }
    return it->second.services[it->second.index++];
}

StubPtr StubManager::getServiceByID(ServiceID id) {
    auto it = serviceMap_.find(id);
    if (it == serviceMap_.end()) {
        return nullptr;
    }
    return it->second;
}

void StubManager::update() {
    auto tick_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
        std::chrono::system_clock::now().time_since_epoch()).count();
    while (timer_.runOnce()) {}
    for (auto& it : serviceMap_) {
        it.second->update(tick_ms);
    }
}

StubManager::ServiceMap& StubManager::getMutableServiceMapUnsafe() {
    return serviceMap_;
}

void StubManager::resetAllStub(ServiceUUID uuid, BundleEntryFunc func) {
    auto it = typeMap_.find(uuid);
    if (it == typeMap_.end()) {
        return;
    }
    for (auto& stub : it->second.services) {
        stub->reset(func, &RpcRef);
    }
}

void StubCallManager::add(StubCallPtr call) {
    callMap_.insert(std::make_pair(callID_, call));
    call->setStubCallID(callID_++);
}

StubCallPtr StubCallManager::get(StubCallID id) {
    auto it = callMap_.find(id);
    if (it == callMap_.end()) {
        return nullptr;
    }
    return it->second;
}

void StubCallManager::destroy(StubCallID id) {
    callMap_.erase(id);
}

StubCreator::StubCreator() {
}

StubCreator::~StubCreator() {
}

StubPtr StubCreator::createService(ServiceUUID uuid, kratos::service::ServiceContext* ctx) {
    auto it = factoryMap_.find(uuid);
    if (it == factoryMap_.end()) {
        return nullptr;
    }
    auto* stub = it->second->create(ctx);
    if (!stub) {
        return StubManagerRef.getServiceByUUID(uuid);
    }
    auto ptr = StubManagerRef.addService(std::shared_ptr<Stub>(stub));
    if (!ptr) {
        return nullptr;
    }
    return ptr;
}

StubPtr StubCreator::getService(ServiceID serviceID) {
    return StubManagerRef.getServiceByID(serviceID);
}

bool StubCreator::isExist(ServiceUUID uuid) {
    return (factoryMap_.find(uuid) != factoryMap_.end());
}

bool StubFactoryDynamic::load(const std::string& bundlePath) {
    for (auto& ci : classList_) {
        if (ci.bundlePath == bundlePath) {
            Logbegin << "Bundle already loaded [" << bundlePath << "]" << Logend;
            return true;
        }
    }
    ClassInfo ci;
    ci.bundlePath = bundlePath;
#ifdef WIN32
    ci.loadHandle = ::LoadLibraryA(bundlePath.c_str());
#else
    ci.loadHandle = dlopen(bundlePath.c_str(), RTLD_LAZY);
#endif // WIN32
    if (ci.loadHandle == INVALID_MODULE_HANDLE) {
        Logbegin << "Load bundle failed [" << bundlePath << "]" << Logend;
        return false;
    }
#ifdef WIN32
    auto regFunc = (BundleRegisterFunc)::GetProcAddress(ci.loadHandle, DEFAULT_REGISTER_NAME);
#else
    auto regFunc = (BundleRegisterFunc)dlsym(ci.loadHandle, DEFAULT_REGISTER_NAME);
#endif
    if (regFunc == nullptr) {
        Logbegin << "Load bundle failed [" << bundlePath << "]" << Logend;
        Logbegin << "Cannot find register entry function [" << DEFAULT_REGISTER_NAME << "]" << Logend;
        unload(ci);
        return false;
    }
    auto uuid = regFunc(RpcRef.getLogger());
    if (!StubCreatorRef.isExist(uuid)) {
        unload(ci);
        Logbegin << "Bundle UUID is not exist in creator " << "[" << uuid << "][" << bundlePath << "]" << Logend;
        return false;
    }
#ifdef WIN32
    ci.entryFunc = (BundleEntryFunc)::GetProcAddress(ci.loadHandle, ci.entryName.c_str());
#else
    ci.entryFunc = (BundleEntryFunc)dlsym(ci.loadHandle, ci.entryName.c_str());
#endif
    if (!ci.entryFunc) {
        unload(ci);
        Logbegin << "Entry function not found [" << bundlePath << "]" << Logend;
        return false;
    }
    classList_.push_back(curClass_);
    curClass_ = ci;
    // Now updating
    onStatusChange(StubStatus::UPDATING);
    return true;
}

void StubFactoryDynamic::unloadClass(bool all) {
    for (auto& ci : classList_) {
        unload(ci);
    }
    if (all) {
        unload(curClass_);
    }
    classList_.clear();
}

bool StubFactoryDynamic::unload(ClassInfo& ci) {
    if (ci.loadHandle == INVALID_MODULE_HANDLE) {
        return true;
    }
#ifdef WIN32
    if (FALSE == ::FreeLibrary(ci.loadHandle)) {
        Logbegin << "System function error, FreeLibrary" << Logend;
        if (FALSE == ::UnmapViewOfFile(ci.loadHandle)) {
            Logbegin << "System function error, UnmapViewOfFile" << Logend;
            return false;
        }
    }
#else
    dlclose(ci.loadHandle);
#endif // WIN32
    ci.loadHandle = INVALID_MODULE_HANDLE;
    ci.entryFunc = nullptr;
    Logbegin << "Unload bundle [" << ci.bundlePath << "]" << Logend;
    return true;
}

bool StubCreator::loadClass(ServiceUUID uuid, const std::string& bundlePath, kratos::service::ServiceContext* ctx) {
    auto it = factoryMap_.find(uuid);
    if (it == factoryMap_.end()) {
        return false;
    }
    if (it->second->isDynamic()) {
        auto* dynamicFactory = dynamic_cast<StubFactoryDynamic*>(it->second.get());
        if (!dynamicFactory->load(bundlePath)) {
            Logbegin << "Load bundle from disk failed [" << bundlePath << "]" << Logend;
            return false;
        }
        if (!createService(uuid, ctx)) {
            Logbegin << "Fork bundle stub failed [" << bundlePath << "]" << Logend;
            return false;
        }
        return true;
    } else {
        Logbegin << "Try to load bundle from disk but stub is static ["
            << bundlePath << "]" << Logend;
        return false;
    }
}

bool StubCreator::unloadClass(ServiceUUID uuid) {
    auto it = factoryMap_.find(uuid);
    if (it == factoryMap_.end()) {
        return false;
    }
    if (it->second->isDynamic()) {
        auto* dynamicFactory = dynamic_cast<StubFactoryDynamic*>(it->second.get());
        dynamicFactory->unloadClass(true);
    } else {
        return false;
    }
    return true;
}

}
